from com.arcfox.middleware import async_mysql_middleware as db
from com.arcfox.manager.redis_task_manager import RedisTaskManager
from loguru import logger


class TaskManager:
    def __init__(self):
        self.manager = RedisTaskManager("yiqing_task")

    async def pull_task(self):
        return await self.manager.pull_tasks(100)

    async def add_tasks(self):
        task_size = await self.manager.task_size()
        if task_size > 0:
            logger.info("队列中还有未消费的任务!")
            return
        sql = f"select city_code from af_city_info"
        rows = await db.query_all(sql)
        tasks = []
        for row in rows:
            tasks.append(row[0])
        await self.manager.add_tasks(tasks)

    async def add_task(self, tasks=[]):
        await self.manager.add_tasks(tasks)
